from sage.crypto.util import bin_to_ascii, ascii_to_bin

class lib:
	def prime_gen(self, n, t):
		kmin = ceil ((2^(n - 1)) / 2)
		kmax = floor((2^n - 2  ) / 2) 
		while 1:
			k = ZZ.random_element(kmin, kmax)
			p = 2*k + 1
			if self.ppt(p, t):  break
		return p

	def ppt(self, n, t):
		v = 1
		w = n - 1
		while is_even(w):
			v = v + 1
			w = ZZ(w / 2)
		for j in range(t):
			a = ZZ.random_element(1, n)
			b = power_mod(a, w, n)
			if((b == 1) or (b == n - 1)):  continue
			for i in range(v):
				b = power_mod(b, 2, n)
				if(b == n - 1):  break
				if(b == 1    ):  return False
			if(b == n - 1): continue
			return False
		return True

	def primitive_root(self, sp):
		p = ZZ((sp - 1) / 2)
		while 1:
			k = self.ZZsecurerandom(sp - 1, 1)
			if(power_mod(k, 2, sp) == 1):  continue
			if(power_mod(k, p, sp) == 1):  continue
			break
		return k

	def int2bin(self, n):
		b = bin(n).lstrip("0")[1:]
		pad = '0' * ((8 - len(b)) % 8)
		return pad + b

	def int2ascii(self, n, bytes = -1):
		b = self.int2bin(n)
		s = bin_to_ascii(b)
		if(bytes == -1   ):  return s
		if(bytes < len(s)):  raise NameError("lib: int2ascii: " + str(n) + " ocupies more than " + str(bytes) + " bytes")
		return '\x00'*(bytes - len(s)) + s

	def ascii2int(self, s):
		b = ascii_to_bin(s)
		n = int(str(b), 2)
		return n

	def mgf(self, hashf, seed, noctets):
		T = ''
		for i in range( ceil(noctets / hashf().digest_size) ):
			counter = self.int2ascii(i)
			T += hashf(seed + counter).digest()
		return T[:noctets]

	def strxor(self, s1, s2):
		if(len(s1) > len(s2)):
			aux = s2
			s2  = s1
			s1  = aux
		s1 = '\x00'*(len(s2) - len(s1)) + s1
		xor = [ord(a).__xor__(ord(b))  for (a,b) in zip(s1, s2)]
		return ''.join( map(chr, xor) )

	def ZZsecurerandom(self, max, min = 0):
		if(max < min):  raise NameError("lib: ZZsecurerandom: max < min")
		while 1:
			nbits  = floor((log(max, 2)) + 1)
			nbytes = ceil(nbits / 8)
			rbytes = os.urandom(nbytes)
			rbits  = str(ascii_to_bin(rbytes))[:nbits]
			r = int(rbits, 2)
			if(r >= min):  break
		return r

	def ZZsecurerandom_nbits(self, nbits):
		nbytes = ceil((nbits - 1) / 8)
		rbytes = os.urandom(nbytes)
		rbits  = '1' + str(ascii_to_bin(rbytes))[:nbits-1]
		return int(rbits, 2)